#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright (C) 2023 reinovo, Inc. All Rights Reserved 
#
# @Time    : 2023/12/29 上午11:33
# @Author  : hmm
# @Email   : liuyuhang0531@foxmail.com
# @File    : transform.py
# @Software: PyCharm

import cv2 as cv
import numpy as np
import apriltag
from math import *
from arm_control import uarm_control

class Transform(object):

    def __init__(self,mtx):
        self.at_detector = apriltag.Detector(apriltag.DetectorOptions(families='tag36h11 tag36h10'))
        self.mtx = mtx

    def eulerAngles2RotationMatrix(self,theta):
        """
        欧拉角转换旋转矩阵
        :param theta: 欧拉角角度
        :return: 旋转矩阵
        """
        R_x = np.array([[1, 0, 0],
                        [0, cos(theta[0]), sin(theta[0])],
                        [0, -sin(theta[0]), cos(theta[0])]
                        ])
        R_y = np.array([[cos(theta[1]), 0, sin(theta[1])],
                        [0, 1, 0],
                        [-sin(theta[1]), 0, cos(theta[1])]
                        ])
        R_z = np.array([[cos(theta[2]), -sin(theta[2]), 0],
                        [sin(theta[2]), cos(theta[2]), 0],
                        [0, 0, 1]
                        ])
        R = R_z @ R_y @ R_x
        return R

    def rotationMatrixToEulerAngles(self,R):
        """
        旋转矩阵转换欧拉角
        :param R: 旋转矩阵
        :return: 欧拉角
        """
        sy = sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
        singular = sy < 1e-6
        if not singular:
            x = atan2(R[2, 1], R[2, 2])
            y = atan2(-R[2, 0], sy)
            z = atan2(R[1, 0], R[0, 0])
        else:
            x = atan2(-R[1, 2], R[1, 1])
            y = atan2(-R[2, 0], sy)
            z = 0
        return np.array([x, y, z])

    def get_img2tag(self, image, tag_size):
        img2tag_p = dict()
        dst = list()
        gray = cv.cvtColor(image,cv.COLOR_BGR2GRAY)
        tags = self.at_detector.detect(gray)
        cam_params = [self.mtx[0, 0], self.mtx[1, 1], self.mtx[0, 2], self.mtx[1, 2]]
        gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
        for tag in tags:
            M, e1, e2 = self.at_detector.detection_pose(tag, cam_params, tag_size)
            angle = self.rotationMatrixToEulerAngles(M[:3, :3]) * 180 / pi
            img2tag_p = {"tag_id":tag.tag_id,"center":tag.center,"M":M,"angle":angle}
            dst.append(img2tag_p)
        return dst

    def get_base2tag_p(self, tag, hand2eye_R, hand2eye_t, base2end_R, base2end_t):
        pix2img_t = np.linalg.inv(self.mtx) @ np.insert(tag["center"], 2, 1)
        img2tag_t = (tag["M"] @ np.insert(pix2img_t, 3, 1))[:3]
        dst = base2end_t + base2end_R @ hand2eye_R @ img2tag_t + base2end_R @ -hand2eye_R @ hand2eye_t
        return dst




